#!/usr/bin/env python

# This example script merges memory or instruction trace files produced by the ARM Instruction
# Emulator memory profiler prototype and DynamoRIO.

import sys
import argparse
import gzip

def merge_memtrace(scalarfname, svefname, outfile, zipped):
    if zipped:
        scalarfile = gzip.open(scalarfname, 'r')
        svefile = gzip.open(svefname, 'r')
    else:
        scalarfile = open(scalarfname, 'r')
        svefile = open(svefname, 'r')

    scalarline = scalarfile.readline()
    if scalarline.startswith("Format"):
        scalarline = scalarfile.readline()
    sveline = svefile.readline()

    while scalarline and sveline:
        scalar_seq = int(scalarline.split(' ')[0][:-1])
        sve_seq = int(sveline.split(' ')[0][:-1])
        assert scalar_seq != sve_seq
        if scalar_seq < sve_seq:
            outfile.write(scalarline)
            scalarline = scalarfile.readline()
        else:
            outfile.write(sveline)
            sveline = svefile.readline()

    if scalarline:
        while scalarline:
            outfile.write(scalarline)
            scalarline = scalarfile.readline()
    if sveline:
        while sveline:
            outfile.write(sveline)
            sveline = svefile.readline()
    
    scalarfile.close()
    svefile.close()
    outfile.close()

def merge_instrace(scalarfname, svefname, outfile, zipped):
    if zipped:
        scalarfile = gzip.open(scalarfname, 'r')
        svefile = gzip.open(svefname, 'r')
    else:
        scalarfile = open(scalarfname, 'r')
        svefile = open(svefname, 'r')

    # Ignore Header in both files (if existent)
    scalarline = scalarfile.readline()
    if scalarline.startswith("Format"):
        scalarline = scalarfile.readline()
    sveline = svefile.readline()
    if sveline.startswith("Format"):
        sveline = sveline.readline()

    while scalarline and sveline:
        scalar_seq = int(scalarline.split(' ')[0][:-1])
        sve_seq = int(sveline.split(' ')[0][:-1])
        scalar_pc =  scalarline.split(' ')[2][:-1]
        sve_pc = sveline.split(' ')[2][:-1]
        assert scalar_seq != sve_seq
        
        # if both traces have a SVE entry, ignore the one in the scalar trace to avoid duplicates
        if scalar_pc == sve_pc and abs(scalar_seq-sve_seq) == 1:
            outfile.write(sveline)
            scalarline = scalarfile.readline()
            sveline = svefile.readline()
        else:
            if scalar_seq < sve_seq:
                outfile.write(scalarline)
                scalarline = scalarfile.readline()
            else:
                outfile.write(sveline)
                sveline = svefile.readline()

    if scalarline:
        while scalarline:
            outfile.write(scalarline)
            scalarline = scalarfile.readline()
    if sveline:
        while sveline:
            outfile.write(sveline)
            sveline = svefile.readline()
    
    scalarfile.close()
    svefile.close()
    outfile.close()

def main():
    parser = argparse.ArgumentParser(description='This script merges memory and instruction traces \
        generated by the Arm Instruction Emulator (ArmIE) and DynamoRIO.\n \
        Usage: --memtrace/instrace <aarch64 scalar instr trace> <SVE vector instr trace>')
    parser.add_argument('--memtrace', action='store_true', help='Merge memtrace files')
    parser.add_argument('--instrace', action='store_true', help='Merge instrace files')
    parser.add_argument('-o', '--out', type=str, help='output file. Default = stdout')
    parser.add_argument('aarch64_trace_file', type=str, help='Aarch64 DR trace')
    parser.add_argument('sve_trace_file', type=str, help='SVE ArmIE scalar trace')
    parser.add_argument('-z', '--zipped', action='store_true', help='Merge zipped trace inputs')

    args = parser.parse_args()
    scalar_file = args.aarch64_trace_file
    sve_file = args.sve_trace_file

    if args.memtrace and args.instrace:
        print "usage: " + sys.argv[0] + "--memtrace/instrace <aarch64 scalar instr trace> <SVE vector instr trace>"
        sys.exit(0)

    if not args.memtrace and not args.instrace:
        print "usage: " + sys.argv[0] + "--memtrace/instrace <aarch64 scalar instr trace> <SVE vector instr trace>"
        sys.exit(0)

    if not args.out:
    	output_file = sys.stdout
    else:
    	if args.zipped:
    		output_file = gzip.open(args.out, 'w')
    	else:
    		output_file = open(args.out, 'w')
    
    if args.memtrace:
        merge_memtrace(scalar_file, sve_file, output_file, args.zipped)

    if args.instrace:
        merge_instrace(scalar_file, sve_file, output_file, args.zipped)

if __name__ == '__main__':                                                       
    main()

